# ---------- 客户端代码 client.py ----------
"""
安全文件传输客户端
功能：加密本地文件并按自定义名称发送到服务器
协议规范：
1. 发送4字节文件名长度（大端序网络字节序）
2. 发送实际文件名字节数据
3. 发送4字节加密数据长度（大端序）
4. 发送加密文件内容
"""

import socket
from cryptography.fernet import Fernet  # 导入对称加密模块

# region 初始化加密模块
# 加载预共享密钥（需与服务器保持一致）
try:
    with open("secret.key", "rb") as key_file:
        key = key_file.read()  # 读取256位(32字节)密钥
except FileNotFoundError:
    print("\033[31m错误：未找到密钥文件 secret.key\033[0m")
    exit(1)

fernet = Fernet(key)  # 创建Fernet加密器实例


# endregion


def send_file(source_path, save_name, host, port):
    """
    文件加密传输主函数
    参数：
    - source_path: 本地待发送文件路径
    - save_name: 服务器保存的文件名
    - host: 服务器IP地址
    - port: 服务器端口号
    安全特性：
    - 使用AES-CBC加密模式（Fernet默认）
    - 数据完整性验证（HMAC签名）
    """
    client = None  # 初始化socket对象
    try:
        # region 文件读取与加密
        try:
            # 二进制模式读取源文件
            with open(source_path, "rb") as f:
                file_data = f.read()  # 注意：大文件可能消耗较多内存
        except FileNotFoundError:
            print(f"\033[31m错误：源文件 {source_path} 未找到\033[0m")
            return

        # 执行加密（自动生成IV，包含在加密数据中）
        encrypted_data = fernet.encrypt(file_data)
        # endregion

        # region 网络连接建立
        client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        client.connect((host, port))  # 建立TCP连接
        # endregion

        # region 文件名传输协议
        # 将文件名编码为UTF-8字节（限制字符集）
        save_name_bytes = save_name.encode('utf-8')

        # 发送文件名长度头（4字节大端序）
        # 协议设计说明：固定长度头便于服务器预分配内存
        client.send(len(save_name_bytes).to_bytes(4, byteorder='big'))

        # 发送实际文件名内容（确保网络字节流传输）
        client.sendall(save_name_bytes)  # sendall保证完整发送
        # endregion

        # region 加密数据传输协议
        # 发送数据长度头（4字节大端序）
        # 安全考虑：防止服务器被大文件DoS攻击
        client.send(len(encrypted_data).to_bytes(4, byteorder='big'))

        # 分块发送加密数据（实际sendall会处理分块）
        client.sendall(encrypted_data)  # 自动处理TCP分包问题
        print(f"\033[32m[✓] {len(encrypted_data)}字节数据已加密发送\033[0m")
        # endregion

    except Exception as e:
        # 统一异常处理（包括网络错误、编码错误等）
        print(f"\033[31m错误：{str(e)}\033[0m")
    finally:
        # 确保关闭网络连接（重要！）
        if client:
            client.close()


if __name__ == "__main__":
    # 服务器配置（生产环境应通过配置文件设置）
    server_ip = "127.0.0.1"  # 本地环回地址
    server_port = 8080  # 需与服务器端口一致

    # region 用户交互界面
    # 获取源文件名（带输入验证）
    while True:
        src_file = input("请输入处理文件的文件名：").strip()
        if src_file:
            break
        print("\033[33m文件名不能为空\033[0m")

    # 获取目标文件名（带基本验证）
    while True:
        dst_file = input("请输入解密文件名称：").strip()
        if dst_file:
            break
        print("\033[33m文件名不能为空\033[0m")
    # endregion

    # 启动文件传输
    send_file(src_file, dst_file, server_ip, server_port)


